﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace SequenceAssembler
{
	class DeBruijnGraph
	{
		public BackGraph bg { get; private set; }
		public int AbundanceThreshold { get; set; }
		public short MultiplicityThreshold { get; set; }
        private Dictionary<long, List<ReadOffset>> m_partialHashTable;
        private int m_partialHashTableMaxCount = 64 * 1024 * 1024;
        private int m_partialHashTableIndex = 0;
        private int m_BatchCount = 1024 * 1024;

        public DeBruijnGraph(BackGraph backGraph) {
			bg = backGraph;
			AbundanceThreshold = 1;
			MultiplicityThreshold = 1;
            m_partialHashTable = new Dictionary<long, List<ReadOffset>>();
		}

        //public void InsertKmers(char[] readStr, int readId)
        //{
        //    foreach (char ch in readStr) {
        //        if (ch == 'N') {
        //            return; //low quality read
        //        }
        //    }
        //    char[] clip = new char[bg.KmerLen];
        //    for (int i = 0; i < readStr.Length - bg.KmerLen + 1; ++i)
        //    {
        //        Array.Copy(readStr, i, clip, 0, bg.KmerLen);
        //        long id = BaseOperation.kmerToLong(clip);
        //        KmerCell cell = bg.LoadCell<KmerCell>(id);
        //        KmerCell twinCell = null;
        //        if (cell == null) {
        //            cell = new KmerCell(id);
        //            long twinId = bg.TwinId(id);
        //            twinCell = new KmerCell(twinId);
        //        }

        //        cell.AddReadOffset(readId, -i);
        //        bg.SaveCell(cell);
        //        if (twinCell != null) {
        //            bg.SaveCell(twinCell);
        //        }
        //    }
        //}

        public void InsertKmers(string readStr, int readId)
        {
            //char[] clip = new char[bg.KmerLen];
            int len = readStr.Length - bg.KmerLen + 1;
            for (int i = 0; i < len; ++i)
            {
                //Array.Copy(readStr, i, clip, 0, bg.KmerLen);
                //long id = BaseOperation.kmerToLong(clip);
                long id = BaseOperation.kmerToLong(readStr, i, bg.KmerLen);
                List<ReadOffset> res;
                if (m_partialHashTable.TryGetValue(id, out res) == true)
                {
                    ReadOffset ro = new ReadOffset(readId, -i);
                    res.Add(ro);
                }
                else
                {
                    res = new List<ReadOffset>();
                    ReadOffset ro = new ReadOffset(readId, -i);
                    res.Add(ro);
                    m_partialHashTable.Add(id, res);
                    long twinId = bg.TwinId(id);
                    res = new List<ReadOffset>();
                    m_partialHashTable.Add(twinId, res);
                }
            }
            m_partialHashTableIndex += len;
            if (m_partialHashTableIndex >= m_partialHashTableMaxCount)
            {
                SavePartialHashTable();
                ClearPartialHashTable();
            }
        }

        public void InsertRemainKmers()
        {
            if (m_partialHashTableIndex != 0)
            {
                SavePartialHashTable();
                ClearPartialHashTable();
            }
        }

        private void SavePartialHashTable()
        {
            int left = m_partialHashTable.Keys.Count;
            KmerCell[] partialCellList;
            int partialCount;
            if (left > m_BatchCount)
            {
                partialCount = m_BatchCount;
            }
            else
            {
                partialCount = left;
            }
            left -= partialCount;
            int index = 0;
            partialCellList = new KmerCell[partialCount];
            foreach (KeyValuePair<long, List<ReadOffset>> kv in m_partialHashTable)
            {
                KmerCell kc = new KmerCell(kv.Key, kv.Value);
                if (index < partialCount)
                {
                    partialCellList[index] = kc;
                    index += 1;
                }
                else
                {
                    MergeAndSavePartialHashTable(partialCellList);
                    index = 0;
                    if (left > m_BatchCount)
                    {
                        partialCount = m_BatchCount;
                    }
                    else
                    {
                        partialCount = left;
                    }
                    left -= partialCount;
                    partialCellList = new KmerCell[partialCount];
                    partialCellList[index] = kc;
                    index += 1;
                }
            }
            if (partialCount != 0)
            {
                MergeAndSavePartialHashTable(partialCellList);
            }
        }

        private void MergeAndSavePartialHashTable(KmerCell[] cellList)
        {
            long[] ids = new long[cellList.Length];
            for (int i = 0; i < cellList.Length; ++i)
            {
                ids[i] = cellList[i].Id;
            }
            KmerCell[] loadedCellList = bg.LoadCells<KmerCell>(ids);
            for (int i = 0; i < loadedCellList.Length; ++i)
            {
                if (loadedCellList[i] != null)
                {
                    cellList[i] = MergeKmerCell(cellList[i], loadedCellList[i]);
                }
            }
            //Console.WriteLine("save cells");
            bg.SaveCells(cellList);
        }

        private KmerCell MergeKmerCell(KmerCell c1, KmerCell c2)
        {
            if (c1.ReadInfo.Count > c2.ReadInfo.Count)
            {
                c2.ReadInfo.AddRange(c1.ReadInfo);
                return c2;
            }
            else
            {
                c1.ReadInfo.AddRange(c2.ReadInfo);
                return c1;
            }

        }

        private void ClearPartialHashTable()
        {
            m_partialHashTable = new Dictionary<long, List<ReadOffset>>();
            m_partialHashTableIndex = 0;
        }

		protected bool Abundant(KmerCell cell, KmerCell twinCell) {
			return (cell.Coverage + twinCell.Coverage) > AbundanceThreshold;
		}

		public void PruneLowAbundance()
		{
			//copy to avoid synchronization problem
			//since it is just a value list we can just copy construct
			List<long> celllist = new List<long>(bg.CellList);
			foreach (long id in celllist)
			{
				KmerCell cell = bg.LoadCell<KmerCell>(id, StatusEnum.Untrusted);
				KmerCell twinCell = null;
				if (cell == null) { continue; }
				long twinId = bg.TwinId(id);
				twinCell = bg.LoadCell<KmerCell>(twinId);
				if (!Abundant(cell, twinCell)) {
					bg.DeleteCell(cell);
					bg.DeleteCell(twinCell);
				} else {
					bg.SaveCell(cell.TrustedCell());
					bg.SaveCell(cell.TrustedCell());
				}
			}
		}

		/// <summary>
		/// Check both nodes to determine the multiplicity.
		/// Multiplicity is estimated by ReadInfo compatibility.
		/// If strong enough add arc.
		/// Only out arc is considered for each node.
		/// Construct SequenceCell.Data by the last base
		/// </summary>
		public void ConnectTrusted()
		{
			long mask = (1L << (bg.KmerLen << 1)) - 1;
			List<long> celllist = new List<long>(bg.CellList);
			foreach (long id in celllist)
			{
				long twinId = bg.TwinId(id);
				KmerCell cell = bg.LoadCell<KmerCell>(id);
				SequenceCell scell = new SequenceCell(id);
				scell.ReadInfo = cell.ReadInfo;
				scell.Data.push_back((byte) (id & 3));
				scell.Coverage = scell.ReadInfo.Count;

				Base[] bases = {Base.A, Base.C, Base.G, Base.T};
				foreach (Base b in bases) {
					long nid = BaseOperation.neighborId(id, b, mask);
					//if (nid == twinId) { continue; } //do not link twins

					var neighReadInfo = getReadInfo(nid);
					if (neighReadInfo == null) { continue; }
					var twinReadInfo = getReadInfo(bg.TwinId(id));
					var twinNeightReadInfo = getReadInfo(bg.TwinId(nid));
					BuildArc(id, nid, scell.ReadInfo, neighReadInfo, twinNeightReadInfo, twinReadInfo, scell);
				}
				bg.SaveCell(scell);
			}
		}

		private List<ReadOffset> getReadInfo(long id) {
			List<ReadOffset> readInfo = null;
			KmerCell kcell = bg.LoadCell<KmerCell>(id, StatusEnum.Trusted);
			if (kcell != null)
			{
				readInfo = kcell.ReadInfo;
			}
			else
			{
				SequenceCell scell = bg.LoadCell<SequenceCell>(id, StatusEnum.Connected);
				if (scell == null) { return null; }
				readInfo = scell.ReadInfo;
			}
			return readInfo;
		}

		private void BuildArc(long src, long dst, List<ReadOffset> srcReadInfo, List<ReadOffset> dstReadInfo,
			List<ReadOffset> srcReadInfoAdd, List<ReadOffset> dstReadInfoAdd, SequenceCell srcCell)
		{
			short multiplicity = 0;
			int i = 0, j = 0;
			while (i < srcReadInfo.Count && j < dstReadInfo.Count)
			{
				if (srcReadInfo[i].ReadId < dstReadInfo[j].ReadId)
				{
					++i; 
				}
				else if (srcReadInfo[i].ReadId > dstReadInfo[j].ReadId)
				{
					++j;
				}
				else // readid compatible
				{
					if (srcReadInfo[i].Offset < dstReadInfo[j].Offset + 1) {
						++i;
					} else if (srcReadInfo[i].Offset > dstReadInfo[j].Offset + 1) {
						++j;
					} else { // offset compatible
						++multiplicity;
						++i; ++j;
					}
				}
			}

			i = 0; j = 0;
			while (i < srcReadInfoAdd.Count && j < dstReadInfoAdd.Count)
			{
				if (srcReadInfoAdd[i].ReadId < dstReadInfoAdd[j].ReadId)
				{
					++i;
				}
				else if (srcReadInfoAdd[i].ReadId > dstReadInfoAdd[j].ReadId)
				{
					++j;
				}
				else // readid compatible
				{
					if (srcReadInfoAdd[i].Offset < dstReadInfoAdd[j].Offset + 1)
					{
						++i;
					}
					else if (srcReadInfoAdd[i].Offset > dstReadInfoAdd[j].Offset + 1)
					{
						++j;
					}
					else
					{ // offset compatible
						++multiplicity;
						++i; ++j;
					}
				}
			}

			if (Strong(multiplicity)) {
				srcCell.ConnectTo(dst, multiplicity);
			}
		}

		protected bool Strong(short multiplicity) {
			return multiplicity > MultiplicityThreshold;
		}

		public void MergeCoverage() {
			List<long> celllist = new List<long>(bg.CellList);
			foreach (long id in celllist) {
				SequenceCell cell = bg.LoadCell<SequenceCell>(id, StatusEnum.Connected);
				if (cell == null) { continue; }
				SequenceCell twinCell = bg.LoadCell<SequenceCell>(bg.TwinId(id));
				float coverage = cell.Coverage + twinCell.Coverage;
				cell.Coverage = coverage;
				twinCell.Coverage = coverage;
				bg.SaveCell(cell.CompleteCell());
				bg.SaveCell(twinCell.CompleteCell());
			}
		}

		public int AddCategory(int lowerId, int len, double var)
		{
			return bg.AddCategory(lowerId,	len, var);
		}

		public static void unitTestInsert()
		{
			using (BackGraph bg = new BackGraph(25))
			{
				DeBruijnGraph dbg = new DeBruijnGraph(bg);
				using (Reader reader = new Reader())
				{
					reader.Open("../../../data/test_reads.fa");
					while (reader.Next())
					{
						dbg.InsertKmers(reader.CurrentRead, reader.CurrentID);
					}
				}
				Console.WriteLine("cell count: {0}", bg.CellCount);
			}
		}

		public static void unitTestPrune()
		{
			using (BackGraph bg = new BackGraph(25))
			{
				bg.Load();
				DeBruijnGraph dbg = new DeBruijnGraph(bg);
				Console.WriteLine("pre prune cell count: {0}", bg.CellCount);
				dbg.PruneLowAbundance();
				Console.WriteLine("post prune cell count: {0}", bg.CellCount);
			}
		}

		public static void unitTestConnect()
		{
			using (BackGraph bg = new BackGraph(25))
			{
				bg.Load();
				DeBruijnGraph dbg = new DeBruijnGraph(bg);
				Console.WriteLine("load cell count: {0}", bg.CellCount);
				Console.WriteLine("Connecting k-mers ......");
				dbg.ConnectTrusted();
				Console.WriteLine("save cell count: {0}", bg.CellCount);
			}
		}

		public static void unitTestCoverage()
		{
			using (BackGraph bg = new BackGraph(25))
			{
				bg.Load();
				DeBruijnGraph dbg = new DeBruijnGraph(bg);
				Console.WriteLine("load cell count: {0}", bg.CellCount);
				Console.WriteLine("Merging coverage ......");
				dbg.MergeCoverage();
				Console.WriteLine("save cell count: {0}", bg.CellCount);
			}
		}
	}
}
